"""A text-completion based LLM."""

import logging

from typing_extensions import Unpack

from graphrag.llm.base import BaseLLM
from graphrag.llm.types import (
    CompletionInput,
    CompletionOutput,
    LLMInput,
)

from .ollama_configuration import OllamaConfiguration
from .types import OllamaClientTypes
from .utils import get_completion_llm_args

log = logging.getLogger(__name__)


class OllamaCompletionLLM(BaseLLM[CompletionInput, CompletionOutput]):
    """A text-completion based LLM."""

    _client: OllamaClientTypes
    _configuration: OllamaConfiguration

    def __init__(self, client: OllamaClientTypes, configuration: OllamaConfiguration):
        self.client = client
        self.configuration = configuration

    async def _execute_llm(
            self,
            input: CompletionInput,
            **kwargs: LLMInput
    ) -> CompletionInput | None:
        args = get_completion_llm_args(
            kwargs.get("model_parameters"), self.configuration
        )
        completion = self.client.generate(prompt=input, **args)
        return completion["response"]
